/**
 * 
 */
package edu.umd.clip.lm.ngram;

import java.io.Serializable;
import java.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class Trie<T> implements Serializable {
	private static final long serialVersionUID = 1L;

	private T data;
	private Map<Integer,Trie<T>> subTries;

	public Trie(T data) {
		this.data = data;
	}
	
	public Trie<T> getSubtrie(int key) {
		if (subTries == null) return null;
		return subTries.get(key);
	}

	public void addSubtrie(int key, Trie<T> trie) {
		if (subTries == null) {
			subTries = new HashMap<Integer, Trie<T>>();
		}
		subTries.put(key, trie);
	}

	public Trie<T> findTrie(int key) {
		if (subTries == null) return null;
		return subTries.get(key);
	}
	
	public Trie<T> findTrie(int ctx[]) {
		return findTrie(ctx, 0);
	}

	Trie<T> findTrie(int ctx[], int idx) {
		int t = ctx[idx];
		Trie<T> subTrie = getSubtrie(t);
		if (idx == ctx.length - 1) {
			return subTrie;
		}
		if (subTrie == null) return null;
		return subTrie.findTrie(ctx, idx+1);
	}

	public T getData() {
		return data;
	}

	public Map<Integer, Trie<T>> getSubTries() {
		return subTries;
	}
	
	
	public static interface Callback<T> {
		void call(Trie<T> data);
	}
	
	public void iterate(int level, Callback<T> callback) {
		if (level == 0) {
			callback.call(this);
			return;
		}
		
		if (subTries != null) {
			for(Trie<T> subtrie : subTries.values()) {
				subtrie.iterate(level-1, callback);
			}
		}
	}
}
